# coding = utf-8

'''
从3DIRCADB中进行展示
'''

import os
import SimpleITK as sitk
from PIL import Image
import numpy as np
import cv2
import matplotlib.pyplot as plt

def show_data(case_id, index):
    root_path = "F:\Dataset\Liver\\3Dircadb"
    root_path = os.path.join(root_path, "case_{}".format(str(case_id).zfill(5)))

    liver_and_tumor_path = os.path.join(root_path, "segmentation")
    raw_path = os.path.join(root_path, "imaging")

    ours_root = "F:\predict\\3dircadb\ours"
    hdenseunet_root = "F:\predict\\3dircadb\hdenseunet"
    munet_root = "F:\predict\\3dircadb\munet"
    unet_root = "F:\predict\\3dircadb\\unet"

    ours_root = os.path.join(ours_root, "case_{}\predict_tumor".format(str(case_id).zfill(5)))
    hdenseunet_root = os.path.join(hdenseunet_root, "case_{}\predict_tumor".format(str(case_id).zfill(5)))
    munet_root = os.path.join(munet_root, "case_{}\predict_tumor".format(str(case_id).zfill(5)))
    unet_root = os.path.join(unet_root, "case_{}\predict_tumor".format(str(case_id).zfill(5)))

    for i in range(len(os.listdir(liver_and_tumor_path))):
        liver_and_tumor = os.path.join(liver_and_tumor_path, "{}.npy".format(str(i).zfill(3)))
        liver_and_tumor = np.load(liver_and_tumor)
        liver = np.zeros(liver_and_tumor.shape)
        liver[liver_and_tumor > 0] = 1
        tumor = np.zeros(liver_and_tumor.shape)
        tumor[liver_and_tumor == 2] = 1

        raw_data = os.path.join(raw_path, "{}.npy".format(str(i).zfill(3)))
        raw_data = np.load(raw_data)

        if tumor.sum() > 0 and i == index:
            ours_image = os.path.join(ours_root, "{}.png".format(str(index).zfill(3)))
            hdenseunet_image = os.path.join(hdenseunet_root, "{}.png".format(str(index).zfill(3)))
            munet_image = os.path.join(munet_root, "{}.png".format(str(index).zfill(3)))
            unet_image = os.path.join(unet_root, "{}.png".format(str(index).zfill(3)))

            if os.path.exists(ours_image):
                ours = Image.open(ours_image).convert("L")
                ours = np.array(ours)
                ours[ours > 0] = 1
            else:
                ours = np.zeros(tumor.shape)
            dice_ours = float(2 * (ours * tumor).sum()) / float(ours.sum() + tumor.sum())

            if os.path.exists(hdenseunet_image):
                hdenseunet = Image.open(hdenseunet_image).convert("L")
                hdenseunet = np.array(hdenseunet)
                hdenseunet[hdenseunet > 0] = 1
            else:
                hdenseunet = np.zeros(tumor.shape)
            dice_hdenseunet = float(2 * (hdenseunet * tumor).sum()) / float(
                hdenseunet.sum() + tumor.sum())

            if os.path.exists(munet_image):
                munet = Image.open(munet_image).convert("L")
                munet = np.array(munet)
                munet[munet > 0] = 1
            else:
                munet = np.zeros(tumor.shape)
            dice_munet = float(2 * (munet * tumor).sum()) / float(munet.sum() + tumor.sum())

            if os.path.exists(unet_image):
                unet = Image.open(unet_image).convert("L")
                unet = np.array(unet)
                unet[unet > 0] = 1
            else:
                unet = np.zeros(tumor.shape)
            dice_unet = float(2 * (unet * tumor).sum()) / float(unet.sum() + tumor.sum())

            print(index, round(dice_ours, 2), round(dice_hdenseunet, 2), round(dice_munet, 2), round(dice_unet, 2))

            label = np.zeros(tumor.shape)
            label[liver == 1] = 1
            label[tumor == 1] = 2
            label = (label/2) * 255
            label = label.astype(np.uint8)
            label = cv2.cvtColor(label, cv2.COLOR_GRAY2BGR)

            ours_lable = np.zeros(ours.shape)
            ours_lable = ours_lable.astype(np.uint8)
            ours_lable = cv2.cvtColor(ours_lable, cv2.COLOR_GRAY2BGR)
            ours_lable[ours > 0] = [0, 255, 0]
            for x in range(ours.shape[0]):
                for y in range(ours.shape[1]):
                    if ours[x,y] > 0:
                        ours_lable[x, y] = [int(ours_lable[x,y,0]*0.5 + label[x,y,0]*0.5),
                                            int(ours_lable[x,y,1]*0.5 + label[x,y,1]*0.5),
                                            int(ours_lable[x,y,2]*0.5 + label[x,y,2]*0.5)]
                    else:
                        ours_lable[x,y] = label[x,y]

            hdenseunet_label = np.zeros(ours.shape)
            hdenseunet_label = hdenseunet_label.astype(np.uint8)
            hdenseunet_label = cv2.cvtColor(hdenseunet_label, cv2.COLOR_GRAY2BGR)
            hdenseunet_label[hdenseunet > 0] = [0, 255, 0]
            for x in range(ours.shape[0]):
                for y in range(ours.shape[1]):
                    if hdenseunet[x, y] > 0:
                        hdenseunet_label[x, y] = [int(hdenseunet_label[x, y, 0] * 0.5 + label[x, y, 0] * 0.5),
                                            int(hdenseunet_label[x, y, 1] * 0.5 + label[x, y, 1] * 0.5),
                                            int(hdenseunet_label[x, y, 2] * 0.5 + label[x, y, 2] * 0.5)]
                    else:
                        hdenseunet_label[x, y] = label[x, y]

            munet_label = np.zeros(ours.shape)
            munet_label = munet_label.astype(np.uint8)
            munet_label = cv2.cvtColor(munet_label, cv2.COLOR_GRAY2BGR)
            munet_label[munet > 0] = [0, 255, 0]
            for x in range(ours.shape[0]):
                for y in range(ours.shape[1]):
                    if munet[x, y] > 0:
                        munet_label[x, y] = [int(munet_label[x, y, 0] * 0.5 + label[x, y, 0] * 0.5),
                                                  int(munet_label[x, y, 1] * 0.5 + label[x, y, 1] * 0.5),
                                                  int(munet_label[x, y, 2] * 0.5 + label[x, y, 2] * 0.5)]
                    else:
                        munet_label[x, y] = label[x, y]

            unet_label = np.zeros(ours.shape)
            unet_label = unet_label.astype(np.uint8)
            unet_label = cv2.cvtColor(unet_label, cv2.COLOR_GRAY2BGR)
            unet_label[unet > 0] = [0, 255, 0]
            for x in range(ours.shape[0]):
                for y in range(ours.shape[1]):
                    if unet[x, y] > 0:
                        unet_label[x, y] = [int(unet_label[x, y, 0] * 0.5 + label[x, y, 0] * 0.5),
                                             int(unet_label[x, y, 1] * 0.5 + label[x, y, 1] * 0.5),
                                             int(unet_label[x, y, 2] * 0.5 + label[x, y, 2] * 0.5)]
                    else:
                        unet_label[x, y] = label[x, y]

            plt.subplot(2, 3, 1)
            plt.imshow(raw_data, cmap="gray")
            plt.subplot(2, 3, 2)
            plt.imshow(label)
            plt.subplot(2, 3, 3)
            plt.imshow(ours_lable)
            plt.subplot(2, 3, 4)
            plt.imshow(hdenseunet_label)
            plt.subplot(2, 3, 5)
            plt.imshow(munet_label)
            plt.subplot(2, 3, 6)
            plt.imshow(unet_label)
            plt.show()

            #save
            '''
            plt.imsave("3_raw.png", raw_data, cmap="gray")
            plt.imsave("3_label.png", label)
            plt.imsave("3_ours.png", ours_lable)
            plt.imsave("3_hdenseunet.png", hdenseunet_label)
            plt.imsave("3_munet.png", munet_label)
            plt.imsave("3_unet.png", unet_label)
            '''





if __name__ == '__main__':
    show_data(case_id=18, index=39)